{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PjRuzfwyImeC"
   },
   "source": [
    "# Bringing Your Own Agent to NeMo Agent Toolkit\n",
    "\n",
    "In this notebook, we'll show you how to integrate an existing agent with the NeMo Agent Toolkit (NAT).\n",
    "\n",
    "You'll learn how to wrap agents from other frameworks so they work smoothly with NAT. This lets you take advantage of NAT features like MCP compatibility, observability, optimization, and profiling in your existing agent systems without refactoring your existing code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Table of Contents\n",
    "- [0.0) Setup](#setup)\n",
    "  - [0.1) Prerequisites](#prereqs)\n",
    "  - [0.2) API Keys](#api-keys)\n",
    "  - [0.3) Installing NeMo Agent Toolkit](#installing-nat)\n",
    "- [1.0) Defining an 'Existing' Agent](#defining-existing-agent)\n",
    "- [2.0) Existing Agent Migration](#migration)\n",
    "  - [2.1) Migration Part 1: Transforming Your Existing Agent into a Workflow](#migration-part-1)\n",
    "  - [2.2) Migration Part 2: Making Your Agent Configurable](#migration-part-2)\n",
    "  - [2.3) Migration Part 3: Integration with NeMo Agent Toolkit](#migration-part-3)\n",
    "  - [2.4) Migration Part 4: A Zero-Code Configuration](#migration-part-4)\n",
    "- [3) Next Steps](#next-steps)\n",
    "\n",
    "<span style=\"color:rgb(0, 31, 153); font-style: italic;\">Note: In Google Colab use the Table of Contents tab to navigate.</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"setup\"></a>\n",
    "# 0.0) Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"prereqs\"></a>\n",
    "## 0.1) Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "msNOf12FDbK7"
   },
   "source": [
    "- **Platform:** Linux, macOS, or Windows\n",
    "- **Python:** version 3.11, 3.12, or 3.13\n",
    "- **Python Packages:** `pip`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6bqpdfFUDdOY"
   },
   "source": [
    "<a id=\"api-keys\"></a>\n",
    "## 0.2) API Keys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i4OTdB6wTdRZ"
   },
   "source": [
    "For this notebook, you will need the following API keys to run all examples end-to-end:\n",
    "\n",
    "- **NVIDIA Build:** You can obtain an NVIDIA Build API Key by creating an [NVIDIA Build](https://build.nvidia.com) account and generating a key at https://build.nvidia.com/settings/api-keys\n",
    "- **Tavily:** You can obtain a Tavily API Key by creating a [Tavily](https://www.tavily.com/) account and generating a key at https://app.tavily.com/home\n",
    "\n",
    "Then you can run the cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if \"NVIDIA_API_KEY\" not in os.environ:\n",
    "    nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key\n",
    "\n",
    "if \"TAVILY_API_KEY\" not in os.environ:\n",
    "    tavily_api_key = getpass.getpass(\"Enter your Tavily API key: \")\n",
    "    os.environ[\"TAVILY_API_KEY\"] = tavily_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RZ3yqbt4CpHH"
   },
   "source": [
    "<a id=\"installing-nat\"></a>\n",
    "## 0.3) Installing NeMo Agent Toolkit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OSICVNHGGm9l"
   },
   "source": [
    "The recommended way to install NAT is through `pip` or `uv pip`.\n",
    "\n",
    "First, we will install `uv` which offers parallel downloads and faster dependency resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install uv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EBV2Gh9NIC8R"
   },
   "source": [
    "NeMo Agent toolkit can be installed through the PyPI `nvidia-nat` package.\n",
    "\n",
    "There are several optional subpackages available for NAT. The `LangChain` subpackage contains useful components for integrating and running within [LangChain](https://python.langchain.com/docs/introduction/). Since LangChain will be used later in this notebook, let's install NAT with the optional `langchain` subpackage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "uv pip show -q \"nvidia-nat-langchain\"\n",
    "if [ $? -ne 0 ]; then\n",
    "    uv pip install \"nvidia-nat[langchain]\"\n",
    "else\n",
    "    echo \"nvidia-nat[langchain] is already installed\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5Qel98PyXOf0"
   },
   "source": [
    "<a id=\"defining-existing-agent\"></a>\n",
    "# 1.0) Defining an 'Existing' Agent\n",
    "\n",
    "In this case study, we will use a simple, self-contained LangChain agent as a proxy for your 'existing' agent. This agent comes equipped with a search tool that is capable of retrieving context from the internet using the Tavily API. The cell below defines the simple LangChain agent with a string input query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile langchain_agent.py\n",
    "import os\n",
    "import sys\n",
    "\n",
    "from langchain import hub\n",
    "from langchain.agents import AgentExecutor\n",
    "from langchain.agents import create_react_agent\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "from langchain_tavily import TavilySearch\n",
    "\n",
    "def existing_agent_main():\n",
    "    if len(sys.argv) < 2:\n",
    "        print(\"Usage: python langchain_agent.py \\\"Your question here\\\"\")\n",
    "        sys.exit(1)\n",
    "    user_input = sys.argv[1]\n",
    "\n",
    "    # Initialize a tool to search the web\n",
    "    search = TavilySearch(\n",
    "        max_results=2,\n",
    "        api_key=os.getenv(\"TAVILY_API_KEY\")\n",
    "    )\n",
    "\n",
    "    # Initialize a LLM client\n",
    "    llm = ChatNVIDIA(\n",
    "        model_name=\"meta/llama-3.3-70b-instruct\",\n",
    "        temperature=0.0,\n",
    "        max_completion_tokens=1024,\n",
    "        api_key=os.getenv(\"NVIDIA_API_KEY\")\n",
    "    )\n",
    "\n",
    "    # Use an open source prompt\n",
    "    prompt = hub.pull(\"hwchase17/react-chat\")\n",
    "\n",
    "    # create tools list\n",
    "    tools = [search]\n",
    "\n",
    "    # Initialize a ReAct agent\n",
    "    react_agent = create_react_agent(\n",
    "        llm=llm,\n",
    "        tools=tools,\n",
    "        prompt=prompt,\n",
    "        stop_sequence=[\"\\nObservation\"]\n",
    "    )\n",
    "\n",
    "    # Initialize an agent executor to iterate through reasoning steps\n",
    "    agent_executor = AgentExecutor(\n",
    "        agent=react_agent,\n",
    "        tools=[search],\n",
    "        max_iterations=15,\n",
    "        handle_parsing_errors=True,\n",
    "        verbose=True\n",
    "    )\n",
    "\n",
    "    # Invoke the agent with a user query\n",
    "    response = agent_executor.invoke({\"input\": user_input, \"chat_history\": []})\n",
    "\n",
    "    # Print the response\n",
    "    print(response[\"output\"])\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    existing_agent_main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dbuiuHg1-fcG"
   },
   "source": [
    "There are three main components to this agent:\n",
    "\n",
    "* a web search tool (Tavily)\n",
    "\n",
    "* an LLM (Llama 3.3)\n",
    "\n",
    "* an agent system prompt (obtained from the internet using `langchain.hub`)\n",
    "\n",
    "The agent is constructed from these three components, then an _agent executor_ is created. Finally, we pass the requested input into the executor and get a response back.\n",
    "\n",
    "All of the components in use come from LangGraph/LangChain, but any other framework or example could also work.\n",
    "\n",
    "Next we will run this sample agent to validate that it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python langchain_agent.py \"Who won the last World Cup?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HZSeyYt3GkzC"
   },
   "source": [
    "<a id=\"migration\"></a>\n",
    "# 2.0) Existing Agent Migration\n",
    "\n",
    "<a id=\"migration-part-1\"></a>\n",
    "## 2.1) Migration Part 1: Transforming Your Existing Agent into a Workflow\n",
    "\n",
    "NAT supports users bringing their own agent into the framework. As the primary entrypoint for agent execution is a NAT Workflow. For the first pass at NAT migration we will create a new workflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat workflow create first_agent_attempt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've created a workflow directory for a new agent, we will continue by migrating the agent's functional code into the new workflow. In the next cell, we have adapted the agent code from the `def existing_agent_main()` into a new method `def first_agent_attempt_function()` which encapsulates the exact same functionality, but is decorated and registered for NAT workflow compatibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile first_agent_attempt/src/first_agent_attempt/first_agent_attempt.py\n",
    "import logging\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class FirstAgentAttemptFunctionConfig(FunctionBaseConfig, name=\"first_agent_attempt\"):\n",
    "    pass\n",
    "\n",
    "\n",
    "@register_function(config_type=FirstAgentAttemptFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def first_agent_attempt_function(_config: FirstAgentAttemptFunctionConfig, _builder: Builder):\n",
    "    import os\n",
    "\n",
    "    from langchain import hub\n",
    "    from langchain.agents import AgentExecutor\n",
    "    from langchain.agents import create_react_agent\n",
    "    from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "    from langchain_tavily import TavilySearch\n",
    "\n",
    "    # Initialize a tool to search the web\n",
    "    search = TavilySearch(\n",
    "        max_results=2,\n",
    "        api_key=os.getenv(\"TAVILY_API_KEY\")\n",
    "    )\n",
    "\n",
    "    # Initialize a LLM client\n",
    "    llm = ChatNVIDIA(\n",
    "        model_name=\"meta/llama-3.3-70b-instruct\",\n",
    "        temperature=0.0,\n",
    "        max_completion_tokens=1024,\n",
    "        api_key=os.getenv(\"NVIDIA_API_KEY\")\n",
    "    )\n",
    "\n",
    "    # Use an open source prompt\n",
    "    prompt = hub.pull(\"hwchase17/react-chat\")\n",
    "\n",
    "    # create tools list\n",
    "    tools = [search]\n",
    "\n",
    "    # Initialize a ReAct agent\n",
    "    react_agent = create_react_agent(\n",
    "        llm=llm,\n",
    "        tools=tools,\n",
    "        prompt=prompt,\n",
    "        stop_sequence=[\"\\nObservation\"]\n",
    "    )\n",
    "\n",
    "    # Initialize an agent executor to iterate through reasoning steps\n",
    "    agent_executor = AgentExecutor(\n",
    "        agent=react_agent,\n",
    "        tools=[search],\n",
    "        max_iterations=15,\n",
    "        handle_parsing_errors=True,\n",
    "        verbose=True\n",
    "    )\n",
    "\n",
    "    async def _response_fn(input_message: str) -> str:\n",
    "        response = agent_executor.invoke({\"input\": input_message, \"chat_history\": []})\n",
    "\n",
    "        return response[\"output\"]\n",
    "\n",
    "    yield FunctionInfo.from_fn(_response_fn, description=\"A simple tool capable of basic internet search\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HAC2XqMWcswN"
   },
   "source": [
    "As you can see above, this is almost the exact same code as your 'existing' agent, but has been refactored to fit within a NAT function registration.\n",
    "\n",
    "The only differences are 1) the definition of a closure function `_response_fn` which captures the instantiated agent executor and uses that to invoke the agent and return the response. And 2) the use of the @register_function decorator.\n",
    "\n",
    "We can also simplify the workflow configuration from:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load first_agent_attempt/configs/config.yml\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile first_agent_attempt/configs/config.yml\n",
    "workflow:\n",
    "  _type: first_agent_attempt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q_QGb4ztd16k"
   },
   "source": [
    "Then we can run the new workflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file first_agent_attempt/configs/config.yml --input \"Who won the last World Cup?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bG7-9kfDfAee"
   },
   "source": [
    "This first pass shows how little effort is required to bring an existing agent into NAT. But as we show in the next section, we can also extend this further to offer better configuration!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "byoF_de3G_oQ"
   },
   "source": [
    "<a id=\"migration-part-2\"></a>\n",
    "## 2.2) Migration Part 2: Making Your Agent Configurable\n",
    "\n",
    "For this next part, we will create another workflow, migrate similar functions as shown in Part 1, but make some important parameters configurable for the entire workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat workflow create second_agent_attempt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IoeGuewrHOvi"
   },
   "source": [
    "Then we can update the agent's function.\n",
    "\n",
    "Below, we expand the configuration to include:\n",
    "\n",
    "* the LLM it should use\n",
    "* configurable values for iterations, verbosity, error handling\n",
    "* an optional description\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile second_agent_attempt/src/second_agent_attempt/second_agent_attempt.py\n",
    "import logging\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.component_ref import FunctionRef\n",
    "from nat.data_models.component_ref import LLMRef\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class SecondAgentAttemptFunctionConfig(FunctionBaseConfig, name=\"second_agent_attempt\"):\n",
    "    llm_model_name: str = Field(description=\"LLM name to use\")\n",
    "    max_iterations: int = Field(default=15, description=\"Maximum number of iterations to run the agent\")\n",
    "    handle_parsing_errors: bool = Field(default=True, description=\"Whether to handle parsing errors\")\n",
    "    verbose: bool = Field(default=True, description=\"Whether to print verbose output\")\n",
    "    description: str = Field(default=\"\", description=\"Description of the agent\")\n",
    "\n",
    "\n",
    "@register_function(config_type=SecondAgentAttemptFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def second_agent_attempt_function(config: SecondAgentAttemptFunctionConfig, builder: Builder):\n",
    "    import os\n",
    "\n",
    "    from langchain import hub\n",
    "    from langchain.agents import AgentExecutor\n",
    "    from langchain.agents import create_react_agent\n",
    "    from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "    from langchain_tavily import TavilySearch\n",
    "\n",
    "    # Initialize a tool to search the web\n",
    "    search = TavilySearch(\n",
    "        max_results=2,\n",
    "        api_key=os.getenv(\"TAVILY_API_KEY\")\n",
    "    )\n",
    "\n",
    "    # Initialize a LLM client\n",
    "    llm = ChatNVIDIA(\n",
    "        model_name=config.llm_model_name,\n",
    "        temperature=0.0,\n",
    "        max_completion_tokens=1024,\n",
    "        api_key=os.getenv(\"NVIDIA_API_KEY\")\n",
    "    )\n",
    "\n",
    "    # Use an open source prompt\n",
    "    prompt = hub.pull(\"hwchase17/react-chat\")\n",
    "\n",
    "    # create tools list\n",
    "    tools = [search]\n",
    "\n",
    "    # Initialize a ReAct agent\n",
    "    react_agent = create_react_agent(\n",
    "        llm=llm,\n",
    "        tools=tools,\n",
    "        prompt=prompt,\n",
    "        stop_sequence=[\"\\nObservation\"]\n",
    "    )\n",
    "\n",
    "    # Initialize an agent executor to iterate through reasoning steps\n",
    "    agent_executor = AgentExecutor(\n",
    "        agent=react_agent,\n",
    "        tools=[search],\n",
    "        **config.model_dump(include={\"max_iterations\", \"handle_parsing_errors\", \"verbose\"})\n",
    "    )\n",
    "\n",
    "    async def _response_fn(input_message: str) -> str:\n",
    "        response = agent_executor.invoke({\"input\": input_message, \"chat_history\": []})\n",
    "\n",
    "        return response[\"output\"]\n",
    "\n",
    "    yield FunctionInfo.from_fn(_response_fn, description=config.description)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9Kv9MgwrIl-b"
   },
   "source": [
    "We can then update the configuration file to include the configuration options which previously were embedded into the agent's code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile second_agent_attempt/configs/config.yml\n",
    "workflow:\n",
    "  _type: second_agent_attempt\n",
    "  llm_model_name: meta/llama-3.3-70b-instruct\n",
    "  max_iterations: 15\n",
    "  verbose: false\n",
    "  description: \"A helpful assistant that can search the internet for information\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KjcxFgEXJKpp"
   },
   "source": [
    "We can then run this modified agent to demonstrate the YAML configuration capabilities of NeMo Agent toolkit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file second_agent_attempt/configs/config.yml --input \"Who won the last World Cup?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gqJyRwy8SLxB"
   },
   "source": [
    "<a id=\"migration-part-3\"></a>\n",
    "## 2.3) Migration Part 3: Integration with NeMo Agent Toolkit\n",
    "\n",
    "NeMo Agent toolkit comes with support for various LLM Providers, Frameworks, and additional components.\n",
    "\n",
    "For this last part of migrating an agent, we will adapt the agent to use built-in toolkit components rather than importing directly from LangChain."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7D8yqHbjC6PG"
   },
   "source": [
    "Changes made below:\n",
    "- changing from LLM model name to an LLM _reference_\n",
    "- adapting the code to query NAT for the LLM and Tools to use\n",
    "- switching to the built-in Tavily Search Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat workflow create third_agent_attempt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile third_agent_attempt/src/third_agent_attempt/third_agent_attempt.py\n",
    "import logging\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.component_ref import FunctionRef\n",
    "from nat.data_models.component_ref import LLMRef\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class ThirdAgentAttemptFunctionConfig(FunctionBaseConfig, name=\"third_agent_attempt\"):\n",
    "    tool_names: list[FunctionRef] = Field(default_factory=list, description=\"List of tool names to use\")\n",
    "    llm_name: LLMRef = Field(description=\"LLM name to use\")\n",
    "    max_iterations: int = Field(default=15, description=\"Maximum number of iterations to run the agent\")\n",
    "    handle_parsing_errors: bool = Field(default=True, description=\"Whether to handle parsing errors\")\n",
    "    verbose: bool = Field(default=True, description=\"Whether to print verbose output\")\n",
    "    description: str = Field(default=\"\", description=\"Description of the agent\")\n",
    "\n",
    "# Since our agent relies on Langchain, we must explicitly list the supported framework wrappers.\n",
    "# Otherwise, the toolkit would not know the correct type to return from the builder\n",
    "\n",
    "@register_function(config_type=ThirdAgentAttemptFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def third_agent_attempt_function(config: ThirdAgentAttemptFunctionConfig, builder: Builder):\n",
    "    import os\n",
    "\n",
    "    from langchain import hub\n",
    "    from langchain.agents import AgentExecutor\n",
    "    from langchain.agents import create_react_agent\n",
    "\n",
    "    # Create a list of tools for the agent\n",
    "    tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)\n",
    "\n",
    "    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)\n",
    "\n",
    "    # Use an open source prompt\n",
    "    prompt = hub.pull(\"hwchase17/react-chat\")\n",
    "\n",
    "    # Initialize a ReAct agent\n",
    "    react_agent = create_react_agent(\n",
    "        llm=llm,\n",
    "        tools=tools,\n",
    "        prompt=prompt,\n",
    "        stop_sequence=[\"\\nObservation\"]\n",
    "    )\n",
    "\n",
    "    # Initialize an agent executor to iterate through reasoning steps\n",
    "    agent_executor = AgentExecutor(\n",
    "        agent=react_agent,\n",
    "        tools=tools,\n",
    "        **config.model_dump(include={\"max_iterations\", \"handle_parsing_errors\", \"verbose\"})\n",
    "    )\n",
    "\n",
    "    async def _response_fn(input_message: str) -> str:\n",
    "        response = agent_executor.invoke({\"input\": input_message, \"chat_history\": []})\n",
    "\n",
    "        return response[\"output\"]\n",
    "\n",
    "    yield FunctionInfo.from_fn(_response_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ogE6tvy3hnKw"
   },
   "source": [
    "We can then update the configuration file to include LLM and Function definitions that before were embedded into the agent's code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile third_agent_attempt/configs/config.yml\n",
    "llms:\n",
    "  nim_llm:\n",
    "    _type: nim\n",
    "    model_name: meta/llama-3.3-70b-instruct\n",
    "    temperature: 0.0\n",
    "    max_tokens: 1024\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "functions:\n",
    "  search:\n",
    "    _type: tavily_internet_search\n",
    "    max_results: 2\n",
    "    api_key: $TAVILY_API_KEY\n",
    "\n",
    "workflow:\n",
    "  _type: third_agent_attempt\n",
    "  tool_names: [search]\n",
    "  llm_name: nim_llm\n",
    "  max_iterations: 15\n",
    "  verbose: false\n",
    "  description: \"A helpful assistant that can search the internet for information\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vOuWML8jhwvM"
   },
   "source": [
    "Finally, we can run this modified agent to demonstrate the flexibility and adaptiveness of using NeMo Agent toolkit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file third_agent_attempt/configs/config.yml --input \"Who won the last World Cup?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8co2555JX5aj"
   },
   "source": [
    "<a id=\"migration-part-4\"></a>\n",
    "## 2.4) Migration Part 4: A Zero-Code Configuration\n",
    "\n",
    "Sometimes NeMo Agent toolkit has all of the components you need already. In cases like these, we can rely on zero code additions. The effect of this is being able to **only** specify a configuration file, demonstrating the power of a batteries-included approach."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ci0-gCYnDMw-"
   },
   "source": [
    "The required components for this base example were:\n",
    "- An LLM (NVIDIA NIM-based)\n",
    "- Tavily Internet Search Tool\n",
    "- ReAct Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile search_agent.yml\n",
    "llms:\n",
    "  nim_llm:\n",
    "    _type: nim\n",
    "    model_name: meta/llama-3.3-70b-instruct\n",
    "    temperature: 0.0\n",
    "    max_tokens: 1024\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "functions:\n",
    "  search:\n",
    "    _type: tavily_internet_search\n",
    "    max_results: 2\n",
    "    api_key: $TAVILY_API_KEY\n",
    "\n",
    "workflow:\n",
    "  _type: react_agent\n",
    "  tool_names: [search]\n",
    "  llm_name: nim_llm\n",
    "  verbose: false\n",
    "  description: \"A helpful assistant that can search the internet for information\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file search_agent.yml --input \"Who won the last World Cup?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KZJ-UKXfZzAW"
   },
   "source": [
    "This concludes the \"Bringing Your Own Agent to NeMo Agent toolkit\" notebook.\n",
    "\n",
    "Throughout this notebook, we've demonstrated a complete migration journey from a standalone agent to a fully integrated NeMo Agent Toolkit workflow. In **Part 1**, we started with a basic ReAct agent using raw API calls and manual prompt engineering. In **Part 2**, we refactored the code to use NAT built-in components, replacing custom implementations with standardized LLM and tool abstractions. In **Part 3**, we elevated the architecture by leveraging the NAT ReAct agent implementation, eliminating the need for custom agent logic entirely. Finally, in **Part 4**, we achieved a zero-code solution using only a YAML configuration file, demonstrating the NAT batteries-included philosophy.\n",
    "\n",
    "By migrating existing agents to NeMo Agent Toolkit, you gain access to a unified platform that standardizes how agents are built, evaluated, and deployed. This migration unlocks powerful capabilities: consistent evaluation frameworks for comparing agent performance across different implementations, systematic optimization through shared tooling and best practices, and comprehensive observability with built-in logging and monitoring. Rather than maintaining custom evaluation scripts, optimization pipelines, and monitoring solutions for each agent, NAT provides these capabilities out of the box, allowing you to focus on solving business problems rather than reinventing infrastructure. This unified approach not only accelerates development but also ensures reproducibility and maintainability across your entire agent ecosystem.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MEtrDDQUjSpX"
   },
   "source": [
    "<a id=\"next-steps\"></a>\n",
    "# 3.0) Next Steps\n",
    "\n",
    "The next exercise in this series is `adding_tools_to_agents.ipynb` where you will interactively learn how to create your own tools and agents with NAT."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "unew_312",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
